import re
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import re
import chromadb
from chromadb.utils import embedding_functions
from docx import Document  # 用于读取docx文件
import os
from openai import OpenAI  # 用于调用阿里云百炼API（兼容OpenAI接口）
import PyPDF2
import pdfplumber  # 用于读取PDF文件
from tools.myes import get_es_client, create_health_index_mapping, store_chunk_in_es


def load_employee_file(file_path):
    """
    加载docx格式的员工手册文件，并对文本进行清洗处理

    参数:
        file_path (str): 员工手册docx文件的路径
    返回:
        str: 清洗后的完整文本内容
    """
    # 初始化docx文档对象
    document = Document(file_path)
    all_text = []  # 存储所有段落文本

    # 遍历文档中的每个段落
    for para in document.paragraphs:
        # 去除段落前后的空白字符
        clean_text = para.text.strip()
        if clean_text:  # 只处理非空段落
            # 移除非中英文、数字及常见标点的特殊字符（保留基本标点和空格）
            clean_text = re.sub(r'[^\u4e00-\u9fa5a-zA-Z0-9\s\.,;:!?()""\'\\-]', '', clean_text)
            # 将连续的空格替换为单个空格（规整空格格式）
            clean_text = re.sub(r'\s+', ' ', clean_text)
            all_text.append(clean_text)

    # 使用双换行符合并所有段落
    return '\n\n'.join(all_text)


def load_employee_file_pdf(file_path):
    """
    加载pdf格式的员工手册文件，并对文本进行清洗处理

    参数:
        file_path (str): 员工手册pdf文件的路径
    返回:
        str: 清洗后的完整文本内容
    """
    try:
        all_text = []  # 存储所有页面文本

        # 使用pdfplumber打开PDF文件（相比PyPDF2能更好地处理文本布局）
        with pdfplumber.open(file_path) as pdf:
            # 遍历PDF文件中的所有页面
            for page in pdf.pages:
                # 从当前页面提取文本内容
                text = page.extract_text()
                if text:  # 确保页面有文本内容
                    # 清洗文本：移除非中英文、数字及常见标点的特殊字符
                    text = re.sub(r'[^\u4e00-\u9fa5a-zA-Z0-9\s\.,;:!?()""\'\\-]', '', text)
                    # 规整空格格式：将连续的空格替换为单个空格
                    text = re.sub(r'\s+', ' ', text)
                    # 去除首尾空白字符
                    text = text.strip()
                    if text:  # 只添加非空文本
                        all_text.append(text)

        # 使用双换行符合并所有页面文本
        return '\n\n'.join(all_text)

    except Exception as e:
        print(f"读取PDF文件出错: {e}")
        return ""


def clean_and_preprocess_text(text):
    """
    对文本进行清洗和预处理，去除无用信息，统一格式，提升后续检索精度

    参数:
        text (str): 原始文本内容
    返回:
        str: 清洗和预处理后的文本
    """
    if not text:
        return ""
    
    # 移除多余的空白字符和换行符
    text = re.sub(r'\s+', ' ', text)
    
    # 移除页眉页脚常见的模式（可根据实际情况调整）
    text = re.sub(r'第\s*\d+\s*页\s*共\s*\d+\s*页', '', text)
    
    # 移除只有数字和标点符号的行
    lines = text.split('\n')
    cleaned_lines = []
    for line in lines:
        # 如果一行中大部分都是中文、英文或数字，则保留
        chinese_chars = len(re.findall(r'[\u4e00-\u9fa5]', line))
        english_digits = len(re.findall(r'[a-zA-Z0-9]', line))
        total_chars = len(line.replace(' ', ''))
        
        # 如果有意义的字符占比超过30%，则保留该行
        if total_chars > 0 and (chinese_chars + english_digits) / total_chars > 0.3:
            cleaned_lines.append(line)
    
    text = '\n'.join(cleaned_lines)
    
    # 进一步清洗：移除只包含标点符号的行
    text = re.sub(r'^[\s\W]+$', '', text, flags=re.MULTILINE)
    
    # 去除首尾空白
    text = text.strip()
    
    return text


def split_document_recursive(text, max_chunk_size=500):
    """
    递归地按文档结构拆分文档，确保每个片段语义完整
    
    参数:
        text (str): 待拆分的文档文本
        max_chunk_size (int): 最大块大小（默认500字符）
    返回:
        list[str]: 文档块列表
    """
    # 如果文本长度小于等于最大块大小，直接返回
    if len(text) <= max_chunk_size:
        return [text]
    
    chunks = []
    current_pos = 0
    text_length = len(text)
    
    while current_pos < text_length:
        # 确定当前块的结束位置
        end_pos = min(current_pos + max_chunk_size, text_length)
        
        # 如果不是最后一块，尝试在语义边界处分割
        if end_pos < text_length:
            # 寻找最近的语义分割点（段落结尾、句子结尾等）
            # 查找可能的分割点（从后往前查找）
            split_points = [
                text.rfind('\n\n', current_pos, end_pos),  # 段落结尾
                text.rfind('。', current_pos, end_pos),     # 中文句号
                text.rfind('.', current_pos, end_pos),      # 英文句号
                text.rfind('！', current_pos, end_pos),     # 中文感叹号
                text.rfind('!', current_pos, end_pos),      # 英文感叹号
                text.rfind('？', current_pos, end_pos),     # 中文问号
                text.rfind('?', current_pos, end_pos),      # 英文问号
            ]
            
            # 过滤掉无效的分割点(-1)
            valid_split_points = [p for p in split_points if p > current_pos]
            
            # 如果找到了合适的分割点，则在该点分割
            if valid_split_points:
                split_point = max(valid_split_points)
                # 确保分割点在合理范围内
                if split_point > current_pos + max_chunk_size // 2:
                    end_pos = split_point + 1  # 包含分割符
        
        # 提取当前块
        chunk = text[current_pos:end_pos].strip()
        if chunk:  # 只添加非空块
            chunks.append(chunk)
        
        # 更新当前位置
        current_pos = end_pos
        
        # 跳过空白字符
        while current_pos < text_length and text[current_pos].isspace():
            current_pos += 1
    
    return chunks


def text_embeddings(texts):
    """
    将输入文本（或文本列表）转换为嵌入向量（调用阿里云百炼text-embedding-v4模型）

    参数:
        texts (str或list[str]): 单个文本字符串或文本列表
    返回:
        list[list[float]]: 文本对应的嵌入向量列表（每个向量为float列表）
    """
    # 初始化OpenAI兼容客户端（连接阿里云百炼服务）
    client = OpenAI(
        api_key=os.getenv("DASHSCOPE_API_KEY"),  # 从环境变量获取API密钥（若无环境变量可直接替换为字符串）
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"  # 阿里云百炼API的兼容接口地址
    )

    # 若输入是单个字符串，转换为列表（统一处理格式）
    if isinstance(texts, str):
        texts = [texts]

    # 阿里云百炼text-embedding-v4模型限制：单次批处理最大10条文本
    batch_size = 10
    embeddings = []  # 存储所有文本的嵌入向量

    # 分批处理文本（避免超过模型批处理限制）
    for i in range(0, len(texts), batch_size):
        # 截取当前批次的文本（从i到i+batch_size，不超过总长度）
        batch = texts[i:i + batch_size]
        # 调用嵌入模型生成向量
        completion = client.embeddings.create(
            model="text-embedding-v4",  # 指定使用的嵌入模型
            input=batch,  # 输入当前批次的文本
            dimensions=1024,  # 指定向量维度（v3/v4模型支持）
            encoding_format="float"  # 输出格式为float（而非base64）
        )

        print(f"Processing batch {i + 1}:{completion}")  # 打印批次处理信息

        # 提取当前批次的嵌入向量（每个文本对应一个向量）
        batch_embeddings = [item.embedding for item in completion.data]
        embeddings.extend(batch_embeddings)  # 合并到总向量列表

    return embeddings


def store_embeddings_in_chroma(chunks, embeddings, collection_name="employee_manual"):
    """
    将文本块和对应的嵌入向量存储到Chroma DB中

    参数:
        chunks (list[str]): 文本块列表
        embeddings (list[list[float]]): 对应的嵌入向量列表
        collection_name (str): Chroma集合名称，默认为"employee_manual"
    返回:
        chromadb.Collection: Chroma集合对象
    """
    # 初始化Chroma客户端
    client = chromadb.PersistentClient(path="./chroma_db")
    
    # 创建或获取集合，不指定嵌入函数以避免冲突
    collection = client.get_or_create_collection(name=collection_name)
    
    # 为每个文本块生成唯一ID
    ids = [f"chunk_{i}" for i in range(len(chunks))]
    
    # 准备元数据
    metadatas = [{"source": "employee_manual", "chunk_index": i} for i in range(len(chunks))]
    
    # 将文本块和嵌入向量添加到集合中
    collection.add(
        ids=ids,
        embeddings=embeddings,
        documents=chunks,
        metadatas=metadatas
    )
    
    print(f"成功将 {len(chunks)} 个文本块存储到Chroma DB集合 '{collection_name}' 中")
    return collection


import re


# def extract_metadata_from_chunk(chunk):
#     """从同程旅行文本块中提取核心元数据（简化版）"""
#     metadata = {
#         "entity_type": "未知",  # 实体类型：酒店/景点/航班/套餐
#         "name": "未知名称",  # 实体名称（如"北京希尔顿酒店"）
#         "attr_type": "通用"  # 属性类型（如"价格/设施/开放时间"）
#     }
#     chunk = chunk.strip()  # 清理空格
#
#     # 1. 识别实体类型（优先匹配高频关键词）
#     entity_keywords = [
#         ("酒店", ["酒店", "民宿", "客栈"]),
#         ("景点", ["景点", "景区", "公园", "博物馆"]),
#         ("航班", ["航班", "飞机", "航班号"]),
#         ("套餐", ["套餐", "行程", "路线"])
#     ]
#     for entity, keywords in entity_keywords:
#         if any(kw in chunk for kw in keywords):
#             metadata["entity_type"] = entity
#             break
#
#     # 2. 提取实体名称（简化：取包含实体关键词的短语）
#     if metadata["entity_type"] != "未知":
#         # 匹配"XXX+实体关键词"（如"上海迪士尼乐园"）
#         match = re.search(f'([^，。；]*?{metadata["entity_type"]})', chunk)
#         if match:
#             metadata["name"] = match.group(1)
#         else:
#             # 若无匹配，取前15字作为名称
#             metadata["name"] = chunk[:15] + ("..." if len(chunk) > 15 else "")
#
#     # 3. 识别核心属性类型（按实体类型匹配关键属性）
#     attr_rules = {
#         "酒店": [
#             ("价格", ["价格", "元/晚", "折扣"]),
#             ("设施", ["WiFi", "停车场", "泳池", "早餐"]),
#             ("政策", ["取消", "退款", "退订"]),
#             ("位置", ["地址", "近地铁", "步行"])
#         ],
#         "景点": [
#             ("时间", ["开放时间", "关门", "停止入园"]),
#             ("门票", ["门票", "票价", "优惠"]),
#             ("预约", ["预约", "订票", "提前"]),
#             ("注意", ["禁止", "允许", "须知"])
#         ],
#         "航班": [
#             ("退改", ["改签", "退票", "手续费"]),
#             ("行李", ["托运", "重量", "手提"]),
#             ("时间", ["起飞", "到达", "延误"]),
#             ("登机", ["航站楼", "值机"])
#         ],
#         "套餐": [
#             ("包含", ["包含", "含", "提供"]),
#             ("价格", ["价格", "成人价", "儿童价"]),
#             ("行程", ["每天", "路线", "安排"])
#         ]
#     }
#     # 按实体类型匹配属性
#     if metadata["entity_type"] in attr_rules:
#         for attr, keywords in attr_rules[metadata["entity_type"]]:
#             if any(kw in chunk for kw in keywords):
#                 metadata["attr_type"] = attr
#                 break
#
#     return metadata


# def store_chunks_in_es(chunks, index_name="health_advice"):
#     """
#     将文本块存储到Elasticsearch中
#
#     参数:
#         chunks (list[str]): 文本块列表
#         index_name (str): Elasticsearch索引名称
#     """
#     # 初始化Elasticsearch客户端
#     es_client = get_es_client()
#
#     # 创建索引映射
#     create_health_index_mapping(es_client, index_name)
#
#     # 存储每个文本块
#     for i, chunk in enumerate(chunks):
#         chunk_id = f"chunk_{i}"
#         metadata = extract_metadata_from_chunk(chunk)
#
#         store_chunk_in_es(
#             es_client=es_client,
#             chunk_id=chunk_id,
#             text_content=chunk,
#             disease=metadata["entity_type"],
#             advice_type=metadata.get("attr_type", "通用信息"),
#             index_name=index_name
#         )
#
#     print(f"成功将 {len(chunks)} 个文本块存储到Elasticsearch索引 '{index_name}' 中")


# def min_max_normalize(scores):
#     """
#     对分数进行最小-最大归一化处理
#
#     参数:
#         scores (list): 分数列表
#     返回:
#         list: 归一化后的分数列表
#     """
#     if not scores:
#         return []
#
#     min_score = min(scores)
#     max_score = max(scores)
#
#     # 如果所有分数相同，则返回相同值的列表
#     if min_score == max_score:
#         return [0.5 for _ in scores]
#
#     # 归一化公式: (score - min) / (max - min)
#     normalized_scores = [(score - min_score) / (max_score - min_score) for score in scores]
#     return normalized_scores
#
#
def store_documents_hybrid(chunks, embeddings, collection_name="employee_manual", es_index_name="health_advice"):
    """
    混合存储文档到Chroma DB和Elasticsearch

    参数:
        chunks (list[str]): 文本块列表
        embeddings (list[list[float]]): 对应的嵌入向量列表
        collection_name (str): Chroma集合名称
        es_index_name (str): Elasticsearch索引名称
    """
    # 存储到Chroma DB
    store_embeddings_in_chroma(chunks, embeddings, collection_name)

    # 存储到Elasticsearch
    # store_chunks_in_es(chunks, es_index_name)

    print(f"完成混合存储：{len(chunks)} 个文档块已存储到Chroma DB和Elasticsearch")


if __name__ == '__main__':
    # 添加项目根目录到sys.path
    import sys
    import os
    sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
    
    file_paths = ['12306.docx','铁路客运规章制度一览.docx']
    all_text = ""
    for file_path in file_paths:
        full_path = os.path.join(os.path.dirname(__file__), file_path)
        if not os.path.exists(full_path):
            print(f"文件不存在: {full_path}")
            continue
            
        if file_path.endswith('.pdf'):
            text = load_employee_file_pdf(full_path)
        else:
            text = load_employee_file(full_path)
        
        all_text += text + "\n\n"
    
    # print("原始文本长度:", len(text))
    
    # 数据清洗与预处理
    cleaned_text = clean_and_preprocess_text(all_text)
    print("清洗后文本长度:", len(cleaned_text))
    
    # 文档拆分
    chunks = split_document_recursive(cleaned_text)
    print(f"拆分为 {len(chunks)} 个语义完整的块")
    
    # 显示前几个块作为示例
    for i, chunk in enumerate(chunks[:3]):
        print(f"\n--- 块 {i+1} ---")
        print(chunk)
    
    # 文本向量化
    print("\n开始文本向量化...")
    embeddings = text_embeddings(chunks)
    print(f"生成了 {len(embeddings)} 个向量，每个向量维度: {len(embeddings[0]) if embeddings else 0}")
    
    # 混合存储到Chroma DB和Elasticsearch
    print("\n开始混合存储...")
    store_documents_hybrid(chunks, embeddings)
    print("混合存储完成")
